fix(source-chat): emit SSE keepalives during generation and add stop-… - #1332
fix(source-chat): emit SSE keepalives during generation and add stop-…#1332AugustoSandim wants to merge 19 commits into
Conversation
…streaming button Keep the source chat SSE connection alive while the LLM generates by sending ignored comments every 15 seconds, preventing proxies (including the Next.js rewrite in front of FastAPI) from dropping the idle connection. On the frontend, wire up an AbortController so users can stop an in-flight stream with a new stop button, abort previous requests when sending a new message, and abort on unmount. Add translations for all supported locales and a characterization test for the keepalive behavior.
… persist user message Convert the source chat graph node and streaming endpoint to async so generation can be cancelled when the client disconnects. Persist the user message to the checkpoint up front via aupdate_state so it survives a mid-generation disconnect. Add a HybridSqliteSaver to delegate LangGraph async checkpointer calls to the existing sync SQLite connection. Update characterization tests for the async path and add a disconnect cancellation test.
There was a problem hiding this comment.
All reported issues were addressed across 21 files
Reply with feedback, questions, or to request a fix.
Re-trigger cubic
… superseded streams Check the checkpoint state before persisting the user message in ; skip the up-front append when the same human message is already the unanswered trailing turn, preventing duplicates on retry after a failed generation. In the frontend, guard the streaming cleanup so a superseded send (a newer message replaced the in-flight one) does not clear the newer stream's loading state or trigger a stale refetch. Also show a disabled spinner in the chat composer when streaming has no cancel callback (e.g. notebook chat) instead of a non-functional Stop button. Update characterization tests: make the keepalive test deterministic with an , and add tests for the duplicate-pending-message skip and normal append-after-AI-message behavior.
There was a problem hiding this comment.
All reported issues were addressed across 4 files (changes from recent commits).
Reply with feedback, questions, or to request a fix.
Re-trigger cubic
…-session streams Add a client-generated to and use it (instead of content equality) to skip the up-front user-message append on retry, while still keeping distinct identical messages. Serialize the snapshot → append → invoke sequence per session with an so concurrent requests for the same thread cannot start duplicate generations. Update the frontend to reuse the trailing unanswered human message id when retrying the same content and generate a fresh uuid otherwise, and to clear streaming state when session creation fails. Add ADR-009 documenting the async-to-sync bridge. Add/update characterization and hook tests for message-id dedup, distinct identical messages, and streaming lifecycle.
There was a problem hiding this comment.
1 issue found across 7 files (changes from recent commits).
Confidence score: 4/5
frontend/src/lib/hooks/use-source-chat.tscouples retry/deduplication and message-ID selection to React state, making security-sensitive behavior harder to test and maintain; extract the selection policy into a pure, directly testable function.
Prompt for AI agents (unresolved issues)
Check if these issues are valid — if so, understand the root cause of each and fix them. If appropriate, use sub-agents to investigate and fix each issue separately.
<file name="frontend/src/lib/hooks/use-source-chat.ts">
<violation number="1" location="frontend/src/lib/hooks/use-source-chat.ts:152">
P2: Custom agent: **Security & testability**
The retry/deduplication policy now lives inside `useSourceChat`, where it is coupled to React state instead of being directly testable. Move message-ID selection into a pure service/domain utility and test both retry reuse and distinct identical messages there.</violation>
</file>
Requires human review: Auto-approval blocked because this review re-detected 1 unresolved issue already reported by Cubic.
Tip: Review your code locally with the cubic CLI to iterate faster.
Re-trigger cubic
The test fixture set `model_override: null`, violating `SourceChatSession.model_override?: string` (the interface narrows the base type's `string | null`). And `sendPromise` was read before TypeScript could prove it assigned (assigned only inside an `act()` callback), triggering TS2454. Drop the invalid field and use a definite-assignment assertion. Co-authored-by: Cursor <cursoragent@cursor.com>
05121ca to
852c1c9
Compare
…sage ids - Replace the global per-session `asyncio.Lock` dict with a refcounted `_SessionLock` that evicts its entry once the last stream releases, preventing an unbounded lock table in a long-lived process. - Add characterization test verifying the lock entry is removed after the stream finishes. - Extract `selectMessageId()` to decide whether to reuse the trailing unanswered human id (retry) or generate a fresh uuid (new turn), and send that real id optimistically so the backend `already_pending` check matches the optimistic entry. - Stop filtering `temp-*` ids on send error since optimistic messages now carry their real ids. - Update ADR-009 to include `aget_tuple` in the async-to-sync delegation surface.
There was a problem hiding this comment.
All reported issues were addressed across 27 files
Requires human review: Auto-approval blocked because this review re-detected 1 unresolved issue already reported by Cubic.
Re-trigger cubic
…d lock waits - Add a 1s disconnect poll interval separate from the 15s SSE keepalive so dropped connections stop generation promptly; only emit keepalives when due. - Track whether the session lock was actually acquired and, if the stream is cancelled while waiting, decrement the holder count without releasing an unheld lock. - Serialize concurrent first-time sends into a single session-create promise, avoid duplicate optimistic bubbles on retry, and roll back the optimistic user message if send fails before persistence. - Use a label for the non-cancellable notebook-chat spinner and add translations for all locales. - Update characterization and hook tests for the new poll interval, lock cancellation, and send races.
There was a problem hiding this comment.
All reported issues were addressed across 23 files (changes from recent commits).
Tip: Review your code locally with the cubic CLI to iterate faster.
Re-trigger cubic
…bles on failure, and support non-secure UUID generation - When a user cancels streaming while an auto-create session request is still in flight, adopt the created session id once it resolves and invalidate the sessions list so the next send does not create another empty session. - On send failure, only roll back an optimistic user message added in the current turn; a retry that reuses a persisted trailing human id keeps its bubble visible until the refetch completes. - Add `createMessageId()` with a `Math.random` fallback for contexts where `crypto.randomUUID` is unavailable (e.g. non-secure HTTP) and use it as the default id generator. - Update tests for the new stop-during-create behavior, retry-failure visibility, and UUID fallback; rename the notebook-chat composer test to clarify "without stop support".
There was a problem hiding this comment.
All reported issues were addressed across 28 files
Reply with feedback, questions, or to request a fix.
Re-trigger cubic
…tend stale-state races - Move the per-session lock and pending human-turn persistence policy from `api/routers/source_chat.py` into a new `api/source_chat_service.py` module, exposing `source_chat_turn` so the router delegates serialization and persistence. - Add focused unit tests for the refcounted session lock (serialization, eviction, and cancelled-waiter cleanup) and the message-id-based pending-turn deduplication; update characterization tests to import from the new module. - In the frontend `useSourceChat` hook, use refs and a per-send generation token so stale session snapshots cannot overwrite messages from a newer stream, resolve authoritative state before choosing the message id, and skip session adoption when the abort comes from unmount rather than Stop. - Add tests for authoritative state resolution, stale-refetch suppression, and unmount abort behavior.
There was a problem hiding this comment.
All reported issues were addressed across 9 files (changes from recent commits).
Requires human review: Auto-approval blocked because this review re-detected 1 unresolved issue already reported by Cubic.
Tip: Review your code locally with the cubic CLI to iterate faster.
Re-trigger cubic
…ssages across session switches - Add `getLogSafeErrorMessage()` to strip `Authorization` headers and bound error length before `console.error`, so axios request configs carrying bearer tokens are never logged. - Use the new helper in `useSourceChat` for session creation, hydration, and send failures. - Track which session the shared `messages` list represents via `messagesSessionRef`; only apply stream chunks and optimistic turns when the list still belongs to the originating session, and re-attach the full accumulated answer when the user switches back. - Add tests verifying auth tokens are absent from logs, sends stay on their original session during pre-send hydration switches, and streaming sessions rehydrate correctly after switching away and back.
There was a problem hiding this comment.
All reported issues were addressed across 4 files (changes from recent commits).
Tip: Review your code locally with the cubic CLI to iterate faster.
Re-trigger cubic
before claiming the shared list and restore console spy safely
There was a problem hiding this comment.
All reported issues were addressed across 32 files
Tip: cubic can generate docs of your entire codebase and keep them up to date. Try it here.
Re-trigger cubic
…superseded stream handling - Add `SessionTurnLease` in `api/source_chat_service.py` to acquire the per-session turn lock by polling instead of blocking, so SSE streams can emit keepalive comments and detect client disconnects while queued behind another turn. - Update `stream_source_chat_response` to poll the lease, emit due keepalives before the final response, and release the lease cleanly on cancellation. - Fail sends in `useSourceChat` when pre-send session hydration errors, rather than minting an unverified message id that would bypass backend dedup. - Ensure only the latest send writes shared state: guard `context_indicators` and message ownership against superseded streams, and refetch the persisted checkpoint after cancellation so the pending turn survives. - Add characterization tests for queued keepalives, queued disconnect cleanup, and due keepalive before completion; update hook tests for cancellation unwind, retry id reuse, and hydration failure.
There was a problem hiding this comment.
3 issues found across 5 files (changes from recent commits).
Confidence score: 3/5
frontend/src/lib/hooks/use-source-chat.tsdoes not claim the newly created session’s empty message list, soownsMessages()can reject the optimistic bubble and streamed answer while the session query loads; update this branch to claim the new session’s list.frontend/src/lib/hooks/use-source-chat.test.tsxrestores itsconsole.errorspy only around the lateractand assertions, so an earlierrenderHookorwaitForfailure can leak the spy into other tests; wrap the full setup and wait sequence in guaranteed cleanup.tests/test_chat_routers_characterization.pyassumes exactly twotime.monotonic()reads, making the keepalive test brittle if control flow changes; use a clock stub or sequence that tolerates the expected additional reads.
Prompt for AI agents (unresolved issues)
Check if these issues are valid — if so, understand the root cause of each and fix them. If appropriate, use sub-agents to investigate and fix each issue separately.
<file name="frontend/src/lib/hooks/use-source-chat.ts">
<violation number="1" location="frontend/src/lib/hooks/use-source-chat.ts:297">
P2: When the first message auto-creates a session, this branch never claims the new session’s empty list, so `ownsMessages()` rejects the optimistic bubble and streamed answer while the session query is still loading. Adopt an empty list for `sessionJustCreated` before sending so the first turn is visible immediately and does not depend on query timing.</violation>
</file>
<file name="frontend/src/lib/hooks/use-source-chat.test.tsx">
<violation number="1" location="frontend/src/lib/hooks/use-source-chat.test.tsx:240">
P3: The `console.error` spy is installed before `renderHook`/`waitFor`, but the `try/finally` that restores it only wraps the later `act` and assertions. If the `waitFor(() => expect(result.current.currentSessionId).toBe('session:1'))` assertion fails, the test throws before entering the `try`, the spy is never restored, and `console.error` stays silenced for every test that follows — the exact leak the sibling 'never logs the request config' test was changed to avoid in this same PR. Move the `renderHook`, `waitFor`, and `act` inside the `try` block so the `finally` always restores the spy.</violation>
</file>
<file name="tests/test_chat_routers_characterization.py">
<violation number="1" location="tests/test_chat_routers_characterization.py:716">
P3: The fake clock in `test_stream_source_chat_emits_due_keepalive_before_final_response` hard-codes exactly two `time.monotonic()` reads (`clock = iter([0.0, 100.0])`). It works only because the current control flow reads the clock once for `last_keepalive` initialization and once in the generation loop, and because the session lock is free so the queue phase never runs. Any legitimate change to the streaming flow that adds a third clock read (an extra keepalive check, a second poll iteration, or a disconnect probe) exhausts the iterator and fails the test with an opaque `StopIteration` before any assertion runs. Give the fake clock a non-raising implementation so future control-flow changes produce a meaningful assertion failure instead of a confusing exception.</violation>
</file>
Tip: Review your code locally with the cubic CLI to iterate faster.
Re-trigger cubic
Co-authored-by: cubic-dev-ai[bot] <191113872+cubic-dev-ai[bot]@users.noreply.github.com>
Co-authored-by: cubic-dev-ai[bot] <191113872+cubic-dev-ai[bot]@users.noreply.github.com>
Co-authored-by: cubic-dev-ai[bot] <191113872+cubic-dev-ai[bot]@users.noreply.github.com>
…ot edit A cubic-dev-ai suggested edit duplicated the try/finally in the "fails the send when pre-send hydration errors" test, leaving a dangling second block that referenced `result` out of scope. It broke tsc, the frontend build, and the test run (ReferenceError: result is not defined). Co-Authored-By: Claude <noreply@anthropic.com>
There was a problem hiding this comment.
All reported issues were addressed across 3 files (changes from recent commits).
Tip: Review your code locally with the cubic CLI to iterate faster.
Re-trigger cubic
…nd hydration The session-created guard left `else if (!sessionJustCreated)` dead after the explicit fresh-session branch was added; fold it into a plain `else`. Co-Authored-By: Claude <noreply@anthropic.com>
…ive test The previous counter froze the clock at 100.0 for every call after the first, which left the due-keepalive check at 0 forever if the generation loop ever iterated twice. A monotonically increasing clock keeps the keepalive-before-final-response regression test meaningful under any number of polls. Co-Authored-By: Claude <noreply@anthropic.com>
Summary
The source chat SSE stream (
POST /api/sources/{id}/chat/sessions/{session_id}/messages) had three gaps that together caused the reply to appear only after the user navigated away (see #1327):user_message, then wrote nothing until the LLM finished. Slow local models mean minutes with zero bytes on the wire, so the Next.js rewrite (or any proxy) drops the idle connection.abortControllerRefbut never instantiated it, so there was no way to stop an in-flight stream.asyncio.to_thread, so even after cancel the model kept burning tokens to completion.This PR fixes all three.
What changed
Backend — keepalive + cancellable generation (
api/routers/source_chat.py,open_notebook/graphs/source_chat.py): ping) every 15s while the model generates, so the connection never goes idle.ainvokein a task so generation is genuinely cancellable.request.is_disconnected()and cancel the in-flight task in afinallyblock.aupdate_stateso it survives a mid-generation disconnect (the frontend refetches the checkpoint on cancel/complete).HybridSqliteSaver, aSqliteSaversubclass with async delegates to the existing sync SQLite connection, soainvoke/aupdate_statework without migrating the module-level sync connection.Frontend — stop button (
frontend/src/lib/api/source-chat.ts,hooks/use-source-chat.ts,components/sources/ChatPanel.tsx,page.tsx)AbortSignalthroughsendMessagetofetch.AbortControllerper message, abort the previous one when sending a new message, and abort on unmount.chat.stoptranslation to all 14 locales.Tests
Related Issue
Fixes #1327
Type of Change
How Has This Been Tested?
uv run pytest)Test Details:
uv run pytest tests/— 658 passedruff checkanduv run python -m mypy— clean: pinglines stream during generation, the Stop button cancels the stream, and the reply streams/refetches when it arrives.Design Alignment
Which design principles does this PR support? (See VISION.md)
Explanation:
Generation now runs on an awaitable coroutine instead of a worker thread, so it can be cancelled the moment the client disconnects instead of running to completion.
Checklist
Code Quality
Testing
make rufforruff check . --fixmake lintoruv run python -m mypy .Documentation
Additional Context
Partial text is still not streamed token-by-token —
ainvokereturns the full result on completion. Streaming individual tokens would require switching toastream/astream_events, which is a separate follow-up.Pre-Submission Verification
Before submitting, please verify: